import json
from typing import Dict, Any

from my_utils.llm_client import llm_chat
import re
from json_repair import repair_json


def robust_extract_json(text, type='list') -> Any:
    ''' 从文本中提取json数据（列表或字典） '''
    # 如果是列表类型，则左括号为'['，右括号为']'；如果是字典类型，则左括号为'{'，右括号为'}'
    if type == 'list':
        left_bracket, right_bracket = '[', ']'
    else:
        left_bracket, right_bracket = '{', '}'

    stack = []
    start_index = -1
    for i, c in enumerate(text):
        if c == left_bracket:
            if not stack:
                start_index = i  # 记录起始位置
            stack.append(c)
        elif c == right_bracket:
            if stack:
                stack.pop()
                if not stack:  # 找到匹配的最外层括号
                    try:
                        return json.loads(text[start_index:i + 1])
                    except json.JSONDecodeError:
                        continue  # 继续查找其他可能
    return None


def fix_incomplete_json_output(raw_results):
    """ 修复可能不完整的 JSON 输出，确保以 "\n}\n```" 结尾 """
    modified_results = []
    for result in raw_results:
        if isinstance(result, str):  # 修正：添加了双引号
            # 找到最后一个句号（包括中文句号和英文句号）的位置
            match = re.search(r'[.。](?=[^.。]*$)', result)  # 查找最后一个句号
            if match:
                last_period_index = match.start()
                # 将最后一个句号后的内容替换为 "\n}\n```"
                result = result[:last_period_index + 1] + '"\n}\n```'  # 修正：添加了双引号
            else:
                # 如果没有找到句号，直接添加结尾
                result += '"\n}\n```'  # 修正：添加了双引号
        modified_results.append(result)
    return modified_results


def xml_outline_to_recursive(self, outline_xml: str) -> Dict[str, Any]:
    ''' 将xml大纲解析为递归结构 '''
    # 示例文档
    example = {
        "title": "论文标题",
        "content": "",
        "refs": [],
        "children": [
            {
                "title": "1. 第一章",
                "content": "",
                "refs": [],
                "children": [
                    {
                        "title": "1.1 第一节",
                        "content": "",
                        "refs": [],
                        "children": [
                            {
                                "title": "1.1.1 第一小节",
                                "content": "1.1.1小节内容",
                                "refs": ["ref1", "ref2"]
                            }
                        ]
                    },
                    {
                        "title": "1.2 第二节",
                        "content": "1.2节内容",
                        "refs": ["ref4"]
                    }
                ]
            },
            {
                "title": "2. 第二章",
                "content": "第二章内容",
                "refs": [],
                "children": [
                    {
                        "title": "2.1 第一节",
                        "content": "2.1节内容",
                        "refs": ["ref2"]
                    }
                ]
            }
        ]
    }
    # 将xml大纲解析为对象列表格式（包括章节标题及引用文献的id列表）
    trans_list_prompt = "请你将xml大纲直接转化为递归结构（每层都包括“title”、“content”、“refs”和“children”字段，如示例所示。其中全部content的内容都为空；只有最里层的refs内容有值，其它refs都为空列表；title字段中要包括标题序号，且在序号前正确数量的添加“#”）（注意：直接给我生成好的json字典，无需其它任何内容）"
    messages = [
        {"role": "user",
         "content": f"大纲：'''{outline_xml}'''\n，递归结构示例：'''{str(example)}'''，\n{trans_list_prompt}"}
    ]
    response = llm_chat(messages, temperature=0.01, model_name='GLM_4_FLASH')
    res = response.choices[0].message.content
    outline_recursive_dict = self.robust_extract_json(res, type='dict')
    return outline_recursive_dict


def remove_formulas_and_tables(markdown_text):
    """ 从 Markdown 文本中删除行外公式和 HTML 表格。 """

    # 删除行外公式（$$...$$）
    text_without_formulas = re.sub(r'\$\$.*?\$\$', '', markdown_text, flags=re.DOTALL)

    # 删除 HTML 表格 (<table>...</table>)
    text_without_tables_and_formulas = re.sub(r'<html.*?>.*?</table>.*?</html.*?>', '', text_without_formulas,
                                 flags=re.DOTALL | re.IGNORECASE)

    return text_without_tables_and_formulas


def extract_simple_name(filename):
    """ 从文件名中提取会议简称或期刊名称。 """
    if filename.startswith("Conf_Paper_Meta_Data_"):
        match = re.search(r"Conf_Paper_Meta_Data_([A-Za-z]+)", filename)
        if match:
            return match.group(1)
        else:
            return ""
    elif filename.startswith("Journal_Paper_Meta_Data_"):
        match = re.search(r"Journal_Paper_Meta_Data_(.+?)_with_whole_text", filename)
        if match:
            return match.group(1)
        else:
            return ""
    else:
        return ""


if __name__ == '__main__':
    markdown_input = """
    这是一个包含行外公式的 Markdown 文本：

    $$
    E = mc^2
    $$

    还有一个表格：

    <html id = "id">
        <table>
            <tr>
                <th>Header 1</th>
                <th>Header 2</th>
            </tr>
            <tr>
                <td>Data 1</td>
                <td>Data 2</td>
            </tr>
        </table>
    </html>
    <table >
        <tr>
            <th>Header 1</th>
            <th>Header 2</th>
        </tr>
        <tr>
            <td>Data 1</td>
            <td>Data 2</td>
        </tr>
    </table>

    这是另一段文本。
    """

    cleaned_markdown = remove_formulas_and_tables(markdown_input)
    print(cleaned_markdown)